# knight_hp_detector_demo.py
import os
import cv2
import mss
import time
import numpy as np
import pydirectinput
from Activate_Window import activate_window
from config import (FPS, DURATION, KNIGHT_HP_THRESHOLD,
                    CAPTURE_WINDOW, KNIGHT_HP_WINDOW)

OUTPUT_DIR = "./knight_hp_logs"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 原始画面截图
def capture_frame(sct):
    img = np.array(sct.grab(CAPTURE_WINDOW))    # BGRA (H, W, 4)
    return img[:, :, :3]                        # BGR  (H, W, 3)

# 从原图提取小骑士、BOSS血量图
def extract_hp_bar(frame,window):
    y = window["top"]
    x = window["left"]
    dy = window["height"]
    dx = window["width"]
    return frame[y:y+dy, x:x+dx]

# 全局信息，只读一次
REF_IMG = cv2.imread("../knight_ref_hp.png")
# 识别小骑士血量图为整数数值信息
def detect_knight_hp_change(cmp_img,threshold=KNIGHT_HP_THRESHOLD):
    match = (cmp_img == REF_IMG).all(axis=2)[0]
    diff = np.diff(match.astype(np.int8), prepend=0, append=0)
    starts = np.where(diff == 1)[0]  # True 段开始
    ends = np.where(diff == -1)[0]  # True 段结束
    count = np.sum(ends - starts >= threshold)
    return count

def main():
    activate_window("Hollow Knight")

    pydirectinput.press("w")
    time.sleep(2)
    pydirectinput.press("w")
    time.sleep(1)
    pydirectinput.press("space")
    time.sleep(3)

    with mss.mss() as sct:
        start_time = time.time()
        count = 0
        prev_hp = 10        # 总生命值
        prev_hp_loss = 0    # 本轮受伤值
        while True:
            next_frame_time = start_time + (count + 1) / FPS
            current_time = time.time()
            if current_time < next_frame_time:
                time.sleep(next_frame_time - current_time)

            frame = capture_frame(sct)
            count += 1
            hp_bar = extract_hp_bar(frame,KNIGHT_HP_WINDOW)
            current_hp_loss = detect_knight_hp_change(hp_bar, KNIGHT_HP_THRESHOLD)

            hp_num = prev_hp - (current_hp_loss if current_hp_loss != prev_hp_loss else 0)

            if current_hp_loss != prev_hp_loss or prev_hp_loss == 0:
                prev_hp_loss = current_hp_loss

            # 血量图像仅用于调试，如果确定没有bug可以不保存图像只接收信号
            # hp_str = str(hp_num)
            # filename = os.path.join(OUTPUT_DIR, f"hp_{count:04d}_{hp_str}.png")
            # cv2.imwrite(filename, hp_bar)

            if hp_num == prev_hp:
                reward = 1 # 鼓励不受伤害，学会躲技能
                # print(f"第{count}步 : 角色无伤 {reward}")
            elif hp_num == 0:
                reward = -10 # 角色死亡，采样结束
                print(f"第{count}步 : 角色死亡,受伤{prev_hp}！ {reward}")
                break
            elif hp_num < prev_hp:
                reward = -5 # 受伤惩罚比造成伤害更大，先躲技能再砍人
                print(f"第{count}步 : 角色受伤{prev_hp-hp_num}！ {reward}")
            else:
                reward = -999
                print("血量状态异常")

            prev_hp = hp_num

            if (time.time() - start_time) >= DURATION:
                print(f"{time.time() - start_time:.1f}s : ⏹️  采样结束（未击败 Boss）")
                break

        total_time = time.time() - start_time
        actual_fps = count / total_time  # 帧数 / 总时间

        print(f"\n📊 采样统计:")
        print(f"   目标时长: {DURATION:.1f}s")
        print(f"   实际耗时: {total_time:.2f}s")
        print(f"   总帧数: {count}")
        print(f"   实际平均采样率: {actual_fps:.2f} FPS")

if __name__ == '__main__':
    main()